import os
import time
import datetime
import optparse
from collections import namedtuple
import subprocess


class BlockIOStats(object):
    BlockInfo = namedtuple(
        "BlockInfo",
        [
            "block_name", "nr_requests"
        ]
    )
    BlockStats = namedtuple(
        "BlockStats",
        [
            "read_ios", "read_merges", "read_sectors", "read_ticks",
            "write_ios", "write_merges", "write_sectors", "write_ticks",
            "in_flight", "io_ticks", "time_in_queue"
        ]
    )
    BlockUsage = namedtuple(
        "BlockUsage",
        [
            "block_name", "read_merge_sec", "read_sec", "read_kb_sec",
            "write_merge_sec", "write_sec", "write_kb_sec",
            "avg_request_size", "ave_queue_size",
            "read_wait", "write_wait", "read_write_wait",
            "util_percent", "service_time"
        ]
    )

    @classmethod
    def get_physical_blocks(cls):
        block_names = list()
        cmd = "lsblk |grep disk|awk '{print $1}'"
        exec_code, exec_result = subprocess.getstatusoutput(cmd)
        for block_name in exec_result.splitlines():
            block_names.append(block_name)
        return block_names

    @classmethod
    def get_block_nr_requests(cls, block_name):
        stat_file = os.path.join("/sys/block/", block_name, "queue/nr_requests")
        with open(stat_file, "r") as fh:
            stat_line = fh.read()
            return int(stat_line)

    @classmethod
    def get_block_info_dict(cls):
        block_dict = dict()
        block_names = cls.get_physical_blocks()
        for block_name in block_names:
            stat_file = os.path.join("/sys/block", block_name, "stat")
            if os.path.exists(stat_file):
                nr_requests = cls.get_block_nr_requests(block_name)
                block_info = cls.BlockInfo(block_name=block_name, nr_requests=nr_requests)
                block_dict[block_name] = block_info
        return block_dict

    @classmethod
    def get_block_stats(cls, block_name):
        stat_line = None
        stat_file = os.path.join("/sys/block", block_name, "stat")
        if not os.path.exists(stat_file):
            return None
        with open(stat_file, "r") as fh:
            stat_line = fh.read()
        stat_info = stat_line.split()
        return cls.BlockStats(
            read_ios=int(stat_info[0]), read_merges=int(stat_info[1]), read_sectors=int(stat_info[2]),
            read_ticks=int(stat_info[3]), write_ios=int(stat_info[4]), write_merges=int(stat_info[5]),
            write_sectors=int(stat_info[6]), write_ticks=int(stat_info[7]), in_flight=int(stat_info[8]),
            io_ticks=int(stat_info[9]), time_in_queue=int(stat_info[10])
        )

    @classmethod
    def show_block_usages(cls, block_usage_list, check_times):
        formatter_01 = "{:10s} {:10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s} {:>10s}"
        formatter_02 = "{:10s} {:10s} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} "
        if (check_times % 10) == 0:
            print("")
            print(formatter_01.format(
                "Time", "Device",
                "rrqm/s", "wrqm/s", "r/s", "w/s",
                "rkB/s", "wkB/s", "avgrq-sz", "avgqu-sz",
                "await", "r_await", "w_await",
                "svctm", "%util"
            ))
        for block_usage in block_usage_list:
            print(formatter_02.format(
                datetime.datetime.now().strftime("%H:%M:%S"), block_usage.block_name,
                block_usage.read_merge_sec, block_usage.write_merge_sec, block_usage.read_sec, block_usage.write_sec,
                block_usage.read_kb_sec, block_usage.write_kb_sec, block_usage.avg_request_size,
                block_usage.ave_queue_size,
                block_usage.read_write_wait, block_usage.read_wait, block_usage.write_wait,
                block_usage.service_time, block_usage.util_percent
            ))
        print("")

    @classmethod
    def get_block_utils(cls, block_info: BlockInfo, curr_stats: BlockStats, prev_stats: BlockStats, check_interval):
        read_sectors = curr_stats.read_sectors - prev_stats.read_sectors
        write_sectors = curr_stats.write_sectors - prev_stats.write_sectors
        read_write_sectors = read_sectors + write_sectors

        read_times = curr_stats.read_ios - prev_stats.read_ios
        write_times = curr_stats.write_ios - prev_stats.write_ios
        read_write_times = read_times + write_times

        read_ticks = curr_stats.read_ticks - prev_stats.read_ticks
        write_ticks = curr_stats.write_ticks - prev_stats.write_ticks
        read_write_ticks = read_ticks + write_ticks

        write_merges = curr_stats.write_merges - prev_stats.write_merges
        read_merges = curr_stats.read_merges - prev_stats.read_merges

        read_merge_sec = read_merges / check_interval
        write_merge_sec = write_merges / check_interval

        read_sec = read_times / check_interval
        write_sec = write_times / check_interval
        io_ticks = curr_stats.io_ticks - prev_stats.io_ticks

        read_kb_sec = read_sectors * 512 / 1024 / check_interval
        write_kb_sec = write_sectors * 512 / 1024 / check_interval

        time_in_queue = (curr_stats.time_in_queue - prev_stats.time_in_queue)

        if read_times == 0:
            read_wait = 0
        else:
            read_wait = read_ticks / read_times
        if write_times == 0:
            write_wait = 0
        else:
            write_wait = write_ticks / write_times
        if read_write_times == 0:
            read_write_wait = 0
            service_time = 0
            avg_request_size = 0
        else:
            read_write_wait = read_write_ticks / read_write_times
            service_time = io_ticks / read_write_times
            avg_request_size = read_write_sectors / read_write_times
        # ave_queue_size = curr_stats.in_flight / block_info.nr_requests
        ave_queue_size = time_in_queue / check_interval / 1000
        util_percent = io_ticks * 100 / check_interval / 1000
        return cls.BlockUsage(
            block_name=block_info.block_name,
            read_merge_sec=round(read_merge_sec, 2),
            read_sec=round(read_sec, 2),
            read_kb_sec=round(read_kb_sec, 2),
            write_merge_sec=round(write_merge_sec, 2),
            write_sec=round(write_sec, 2),
            write_kb_sec=round(write_kb_sec, 2),
            avg_request_size=round(avg_request_size, 2),
            ave_queue_size=round(ave_queue_size, 2),
            read_wait=round(read_wait, 2),
            write_wait=round(write_wait, 2),
            read_write_wait=round(read_write_wait, 2),
            service_time=round(service_time, 2),
            util_percent=round(util_percent, 2),
        )

    @classmethod
    def show_io_stats(cls, check_interval=1):
        block_info_dict = cls.get_block_info_dict()
        prev_block_stats_dict = dict()
        check_times = 0
        while True:
            curr_block_stats_dict = dict()
            block_usage_list = list()
            for block_name in block_info_dict.keys():
                curr_block_stats_dict[block_name] = cls.get_block_stats(block_name=block_name)
            if prev_block_stats_dict is not None:
                for block_name in prev_block_stats_dict.keys():
                    block_info = block_info_dict[block_name]
                    curr_stats = curr_block_stats_dict[block_name]
                    prev_stats = prev_block_stats_dict[block_name]
                    block_usage = cls.get_block_utils(
                        block_info=block_info,
                        curr_stats=curr_stats,
                        prev_stats=prev_stats,
                        check_interval=check_interval
                    )
                    block_usage_list.append(block_usage)
            cls.show_block_usages(block_usage_list, check_times)
            time.sleep(check_interval)
            prev_block_stats_dict = curr_block_stats_dict
            check_times += 1
        pass


class CommandHelper(object):

    @classmethod
    def get_command_option(cls):
        usage = """ 获取磁盘使用率,和Linux命令iostat -dxk xx相同。"""
        # 使用使用add_help_option=True或使用默认值，都会自动添加help选项
        # parser = optparse.OptionParser(usage=usage)
        # 如果参数中使用到——h或主动添加--help参数，则使用add_help_option=False 关闭自动添加help选项
        parser = optparse.OptionParser(
            usage=usage,
            add_help_option=True,
            version="iostat 1.0"
        )
        parser.add_option(
            "-i", "--interval",
            action="store",
            type="int",
            dest="check_interval",
            default=1,
            help="检查间隔,默认1秒"
        )
        (options, args) = parser.parse_args()
        return options


def main():
    options = CommandHelper.get_command_option()
    BlockIOStats.show_io_stats(check_interval=options.check_interval)


if __name__ == '__main__':
    main()
